package cn.lgwen.spark.ml.learning.cosmic_radiation;

import cn.lgwen.spark.ml.learning.kaggle.TitanicUtil;
import ml.dmlc.xgboost4j.scala.spark.XGBoostClassificationModel;
import ml.dmlc.xgboost4j.scala.spark.XGBoostClassifier;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.SparkSession;

/**
 * 2020/3/25
 * aven.wu
 * danxieai258@163.com
 */
public class RadiationXGBoostClassifier {

    public static void main(String[] args) {
        SparkSession spark = SparkSession
                .builder().master("local[*]")
                .appName("TitanicRandomForestClass")
                .getOrCreate();

        java.util.Map<String, Object> map = new java.util.HashMap<>();
        map.put("eta", 0.1f);
        map.put("missing", -999);
        map.put("objective", "multi:softprob");
        map.put("num_class", 3);
        map.put("num_round", 100);
        map.put("num_workers", 2);

        XGBoostClassifier classifier = new XGBoostClassifier()
                .setLabelCol("Survived")
                .setFeaturesCol("features");
        Dataset<Row> trainSet = TitanicUtil.trainData(spark);
        XGBoostClassificationModel model = classifier.fit(trainSet);
        Dataset<Row> testSet = TitanicUtil.testData(spark);
        Dataset<Row> res = model.transform(testSet);
        System.out.println(TitanicUtil.evaluate(res));

    }

}
